##__docformat__ = "restructuredtext en"
# ******NOTICE***************
# two_game.py module by Shoibal Chakravarty
#
# You may copy and use this module as you see fit with no
# guarantee implied provided you keep this notice in all copies.
# *****END NOTICE************
"""
This module contains the functions that simulate a scenario where all regions
have a two period planning horizon (current period and the next period, where each region is myopic)
The three regions play a Nash game and the solution is the intersection of the best response surfaces 
which is determined numerically. The green energy deployment in the last period is myopic.
"""

__author__ = "Shoibal Chakravarty (shoibalc@princeton.edu)"
__version__ = "0.1"
__date__ = "2011/August"
__license__ = "BSD"
__all__ = ['threenash','threenashupdate','threeintersectionregion']
import numpy as np
from myopic import  myopicupdate, regionupdate, worldupdate, cumulativecost
from flat2 import flatten
from utility import U3
from two_game import finalperiodupdate
import scipy.optimize

def threenash(r1,r2,r3,r4,w,com,precision=0.5,init_grid_size=40):
    """
    Function threenash solves the scenario where the regions r1, r2 and r3 play a two period
    Nash game (2nd period is myopic). For every period except the last r1, r2 and r3 play their best response
    Nash strategy. The best response of region r1 is a surface that gives its green energy deployment asa function 
    of deployments of r2 and r3. Similarly for r2 and r3. The surfaces  and their intersection are calculated 
    numerically to a given tolerance. In case of multiple solutions, a random solution is picked. The last period
    solution is myopic and is calculated when  penultimate period Nash game is solved.\n\n
    threenash(r1,r2,r3,...):\n
    r1, r2 and r3 play a three region Nash game where they look ahead one period.
    """
    r1.descriptor = r2.descriptor = r3.descriptor = 'Solution of a three region Nash game.'
    r4.descriptor= 'Myopic'
    print "\n\nIn this scenario %s, %s  and %s play a three region two period Nash game\n"% (r1.name, r2.name, r3.name)
    for t in range(w.T-1):
# three region two period nash
        r4.g[t], r4.u1[t], r4.c1[t], r4.m[t] = myopicupdate(r4,r1,r2,r3,w,t)
        regionupdate(r4,t)
        threenashupdate(r1,r2,r3,r4,w,com,t,precision,init_grid_size)
        worldupdate(r1,r2,r3,r4,w,com,t)
        print "\nPeriod", t+1,"\n", r1.name, "g", r1.g[t], "e", r1.e[t],"\n", r2.name, "g", r2.g[t], "e", r2.e[t],\
        "\n",r3.name, "g", r3.g[t], "e", r3.e[t], "\n",r4.name, "g", r4.g[t], "e", r4.e[t], "\ncg", w.cg[t+1],"\n"
# the last period is still myopic
    t = w.T-1
    finalperiodupdate(r1,t)
    finalperiodupdate(r2,t)
    finalperiodupdate(r3,t)
    r4.g[t], r4.u1[t], r4.c1[t], r4.m[t] = myopicupdate(r4,r1,r2,r3,w,t)
    regionupdate(r4,t)
    worldupdate(r1,r2,r3,r4,w,com,t)
    print "\nPeriod", t+1,"\n", r1.name, "g", r1.g[t], "e", r1.e[t],"\n", r2.name, "g", r2.g[t], "e", r2.e[t],\
    "\n",r3.name, "g", r3.g[t], "e", r3.e[t], "\n",r4.name, "g", r4.g[t], "e", r4.e[t], "\ncg", w.cg[t+1],"\n"
    cumulativecost(r1,r2,r3,r4,w)


def threenashupdate(ra,rb,rc,rd,w,com,t,precision,init_grid_size):
    """
    Function threenashupdate solves the three region two period Nash game for the current period (except for 
    the last period in the simulation).  The function successively calls function threeintersectionregion which returns
    a best guess and a cuboid bounding the solution. This procedure is repeated till the desired precision is reached.
    """
    maxa = min(ra.M[t]*ra.E[t+1],(ra.E[t+1] - ra.G[t]))
    mina = 1e-6
    maxb = min(rb.M[t]*rb.E[t+1],(rb.E[t+1] - rb.G[t]))
    minb = 1e-6
    maxc = min(rc.M[t]*rc.E[t+1],(rc.E[t+1] - rc.G[t]))
    minc = 1e-6
    D0 = 19
    D = D0
    EPS = EPSa = EPSb = EPSc = init_grid_size*D
    print 'Search region',ra.name,'min = 0 max =', maxa, rb.name, 'min = 0, max =', maxb, rc.name, 'min = 0, max =',maxc
    flag = True
    num = False
    xbest, ybest, zbest = np.nan, np.nan, np.nan
    while EPS > precision and flag:
# initial run with  init_grid_size. Subsequent runs reduce it by a factor of D
        if num:
            D = np.random.randint(D0+1,D0+5)
        EPSa = EPSa/D
        EPSb = EPSb/D
        EPSc = EPSc/D
        num = True
        print "GRID SIZE", D,
        mina, maxa, minb, maxb, minc, maxc, xbest1, ybest1, zbest1, lenxyz = \
                threeintersectionregion(ra,rb,rc,rd,w,com,t,mina,maxa,minb,maxb,minc,maxc,EPSa,EPSb,EPSc,precision)
        print "Search region update:", ra.name, "min", mina, "max", maxa, "best", xbest1, rb.name, "min", minb, "max", maxb, "best", ybest1, \
                rc.name, "min", minc, "max", maxc, "best", zbest1
        err_best = (xbest - xbest1)**2 + (ybest - ybest1)**2 + (zbest - zbest1)**2
        if np.sqrt(err_best) <= precision:
# this hack prevents infinite looping if we hit a multiple equilibria. we select a particular solution and carry on.
# it also helps in avoiding an extra round of brute force maximization in some cases.
            if lenxyz == 1:
                print "1. ATTAINED REQUIRED PRECISION"
                flag = False
            else:
                print "2. MULTIPLE SOLUTIONS DETECTED, PICKING AVERAGE OF DEGENERATE SOLUTIONS"
                flag = False
            break
        else:
            xbest, ybest, zbest = xbest1, ybest1, zbest1
        EPSa = maxa - mina
        EPSb = maxb - minb
        EPSc = maxc - minc
        EPS = max(EPSa, EPSb, EPSc)
    if flag:
        print "3. ATTAINED REQUIRED PRECISION"
        ga = 0.5*(mina + maxa)
        gb = 0.5*(minb + maxb)
        gc = 0.5*(minc + maxc)
    else:
        ga = xbest
        gb = ybest
        gc = zbest
    
    parama = \
    (gb+gc+rd.g[t], com.cg0, com.greenpar, ra.B[t], ra.B[t+1], ra.cn[t], ra.co[t], ra.cn[t+1], ra.co[t+1], ra.G[t],\
    rb.G[t]+rc.G[t]+rd.G[t], w.G[0], ra.e[t], ra.E[t+1], ra.e[t+1], ra.E[t+2], ra.M[t], ra.M[t+1], (1 - com.discount)**com.years, True)
    u, ra.g[t], ra.u1[t], ra.c1[t], ra.g2[t], ra.u2[t], ra.c2[t], ra.m[t] =  U3(ga,*parama)
    regionupdate(ra,t)

    paramb = \
    (gc+ga+rd.g[t], com.cg0, com.greenpar, rb.B[t], rb.B[t+1], rb.cn[t], rb.co[t], rb.cn[t+1], rb.co[t+1], rb.G[t],\
    ra.G[t]+rc.G[t]+rd.G[t], w.G[0], rb.e[t], rb.E[t+1], rb.e[t+1], rb.E[t+2], rb.M[t], rb.M[t+1], (1 - com.discount)**com.years, True)
    u, rb.g[t], rb.u1[t], rb.c1[t], rb.g2[t], rb.u2[t], rb.c2[t], rb.m[t] =  U3(gb,*paramb)
    regionupdate(rb,t)

    paramc = \
    (ga+gb+rd.g[t], com.cg0, com.greenpar, rc.B[t], rc.B[t+1], rc.cn[t], rc.co[t], rc.cn[t+1], rc.co[t+1], rc.G[t],\
    ra.G[t]+rb.G[t]+rd.G[t], w.G[0], rc.e[t], rc.E[t+1], rc.e[t+1], rc.E[t+2], rc.M[t], rc.M[t+1], (1 - com.discount)**com.years, True)
    u, rc.g[t], rc.u1[t], rc.c1[t], rc.g2[t], rc.u2[t], rc.c2[t], rc.m[t] =  U3(gc,*paramc)
    regionupdate(rc,t)


def threeintersectionregion(r1,r2,r3,r4,w,com,t,mina,maxa,minb,maxb,minc,maxc,EPSa,EPSb,EPSc,precision):
    """
    Function threeintersectionregion calculates the three best response surfaces numerically. It then finds the 
    intersection of these surfaces and returns the best guess and a cuboid bounding the solution. The best response
    surface of region r1 is a function of the sum of the green deployment of r2 and r3. This is used to considerably 
    simplify the calculation of the surface. We first determine:\n
    (r1,x) = (f(x),x)  where f(x) is the best response for a given x.\n\n
    Then r1 and r2 are mapped to x:\n
    r1+f2 -> x\n
    to give the surface:\n
    r1,r2,r3 = (f(x),r1,r2) where x = r1+r2\n\n
    Similarly, for r2 and r3. The intersection of these surfaces (or discretized approximations of these 
    surfaces) are determined by converting these to three dimensional logical sets and 
    finding the points of intersection. 
    
    
    """
#   Ga0,Gb0,Gc0,Gc1,G0,ea0,ea1,Ea0,Ea1,eb0,eb1,Eb0,Eb1,delta,precision,alpha, mina, maxa, minb, maxb, EPSa, EPSb = args

    Wa = 3
    lena = int(np.ceil((maxa - mina)/EPSa) + 1)
    lenb = int(np.ceil((maxb - minb)/EPSb) + 1)
    lenc = int(np.ceil((maxc - minc)/EPSc) + 1)
    lenbc = int(np.ceil((maxb - minb + maxc  - minc)/max(EPSb, EPSc)) + 1)
    lenca = int(np.ceil((maxc - minc + maxa - mina)/max(EPSc, EPSa)) + 1)
    lenab = int(np.ceil((maxa - mina + maxb - minb)/max(EPSa, EPSb)) + 1)
    EPSa = (maxa - mina)/(lena-1)
    EPSb = (maxb - minb)/(lenb-1)
    EPSc = (maxc - minc)/(lenc-1)
    EPSbc = (maxb - minb + maxc  - minc)/(lenbc-1)
    EPSca = (maxc - minc + maxa - mina)/(lenca-1)
    EPSab = (maxa - mina + maxb - minb)/(lenab-1)
    eps = 1e-5
    slice_a = np.s_[mina: maxa + eps:EPSa,]
    slice_b = np.s_[minb: maxb + eps: EPSb,]
    slice_c = np.s_[minc: maxc + eps: EPSc,]
    slice_bc = np.s_[minb + minc: maxb + maxc + eps: EPSbc,]
    slice_ca = np.s_[minc + mina: maxc + maxa + eps: EPSca,]
    slice_ab = np.s_[mina + minb: maxa + maxb + eps:EPSab,]
    mgrid_a = np.mgrid[slice_a[0]]
    mgrid_b = np.mgrid[slice_b[0]]
    mgrid_c = np.mgrid[slice_c[0]]
    mgrid_bc = np.mgrid[slice_bc[0]]
    mgrid_ca = np.mgrid[slice_ca[0]]
    mgrid_ab = np.mgrid[slice_ab[0]]
    lista = np.empty((lenbc+1))
    listb = np.empty((lenca+1))
    listc = np.empty((lenab+1))
    m=0
#    print EPSa,EPSb,EPSc,EPSbc,EPSca,EPSab
    for tt in mgrid_bc:
        parama = \
        (tt+r4.g[t], com.cg0, com.greenpar, r1.B[t], r1.B[t+1], r1.cn[t], r1.co[t], r1.cn[t+1], r1.co[t+1], r1.G[t],\
        r2.G[t]+r3.G[t]+r4.G[t], w.G[0], r1.e[t], r1.E[t+1], r1.e[t+1], r1.E[t+2], r1.M[t], r1.M[t+1], (1 - com.discount)**com.years, False)
        x0 = scipy.optimize.brute(U3,slice_a,args=parama,full_output=False,finish=None)
        lista[m] = x0
        m+=1
    m=0
    for tt in mgrid_ca:
        paramb = \
        (tt+r4.g[t], com.cg0, com.greenpar, r2.B[t], r2.B[t+1], r2.cn[t], r2.co[t], r2.cn[t+1], r2.co[t+1], r2.G[t],\
        r1.G[t]+r3.G[t]+r4.G[t], w.G[0], r2.e[t], r2.E[t+1], r2.e[t+1], r2.E[t+2], r2.M[t], r2.M[t+1], (1 - com.discount)**com.years, False)
        x0 = scipy.optimize.brute(U3,slice_b,args=paramb,full_output=False,finish=None)
        listb[m] = x0
        m+=1
    m=0
    for tt in mgrid_ab:
        paramc = \
        (tt+r4.g[t], com.cg0, com.greenpar, r3.B[t], r3.B[t+1], r3.cn[t], r3.co[t], r3.cn[t+1], r3.co[t+1], r3.G[t],\
        r1.G[t]+r2.G[t]+r4.G[t], w.G[0], r3.e[t], r3.E[t+1], r3.e[t+1], r3.E[t+2], r3.M[t], r3.M[t+1], (1 - com.discount)**com.years, False)
        x0 = scipy.optimize.brute(U3,slice_c,args=paramc,full_output=False,finish=None)
        listc[m] = x0
        m+=1
# MAKE A DICTIONARY OF THE x vs y + z VECTORS TO CONSTRUCT THE x(y,z) SURFACE
    dict_a_bc = dict(zip(range(lenbc),lista))
    dict_b_ca = dict(zip(range(lenca),listb))
    dict_c_ab = dict(zip(range(lenab),listc))
#    print dict_a_bc
#    print dict_b_ca
#    print dict_c_ab
# CONSTRUCT THE INTERSECTION OF THE THREE SURFACES [a(b,c),b,c], [a,b(c,a),c] and [a,b,c(a,b)]
    surface_a = np.empty((lenb,lenc),dtype=int)
    surface_b = np.empty((lena,lenc),dtype=int)
    surface_c = np.empty((lena,lenb),dtype=int)
    vola = False*np.ndarray((lena,lenb,lenc),dtype=bool)
    volb = False*np.ndarray((lena,lenb,lenc),dtype=bool)
    volc = False*np.ndarray((lena,lenb,lenc),dtype=bool)
    for b1 in range(lenb):
        for c1 in range(lenc):
            surface_a[b1,c1] = a1 = int((dict_a_bc[int(round((mgrid_b[b1] + mgrid_c[c1] - minb - minc)/EPSbc))] - mina)/EPSa)
            vola[a1,b1,c1] = True
#    print surface_a
    for a1 in range(lena):
        for c1 in range(lenc):
            surface_b[a1,c1] = b1 = int((dict_b_ca[int(round((mgrid_c[c1] + mgrid_a[a1] -minc - mina)/EPSca))] - minb)/EPSb)
            volb[a1,b1,c1] = True
#    print surface_b
    for a1 in range(lena):
        for b1 in range(lenb):
            surface_c[a1,b1] =c1 = int((dict_c_ab[int(round((mgrid_a[a1] + mgrid_b[b1] -mina -minb)/EPSab))] - minc)/EPSc)
            volc[a1,b1,c1] = True
#    print surface_c
    volc &= vola & volb
    xa,xb,xc = np.nonzero(volc)
#    print xa,xb,xc
    a0 = [mgrid_a[zz] for zz in xa]
    b0 = [mgrid_b[zz] for zz in xb]
    c0 = [mgrid_c[zz] for zz in xc]
    abc0 = zip(a0,b0,c0)
    a_best,b_best,c_best = np.mean(a0),np.mean(b0),np.mean(c0)
    print "CURRENT GUESS(ES)",
    print abc0
    
    xa1 = [xx - Wa > 0 and mgrid_a[xx - Wa].tolist() or mgrid_a[0].tolist() for xx in xa]
    xb1 = [xx - Wa > 0 and mgrid_b[xx - Wa].tolist() or mgrid_b[0].tolist() for xx in xb]
    xc1 = [xx - Wa > 0 and mgrid_c[xx - Wa].tolist() or mgrid_c[0].tolist() for xx in xc]
    xa2 = [xx + Wa < (lena - 1) and mgrid_a[xx + Wa].tolist() or mgrid_a[-1].tolist() for xx in xa]
    xb2 = [xx + Wa < (lenb - 1) and mgrid_b[xx + Wa].tolist() or mgrid_b[-1].tolist() for xx in xb]
    xc2 = [xx + Wa < (lenc - 1) and mgrid_c[xx + Wa].tolist() or mgrid_c[-1].tolist() for xx in xc]
#    print xa1,'\n',xb1,'\n',xc1,'\n',xa2,'\n',xb2,'\n',xc2,'\n'
    mina, maxa, minb, maxb, minc, maxc = min(flatten(xa1)),max(flatten(xa2)),min(flatten(xb1)),max(flatten(xb2)),min(flatten(xc1)),max(flatten(xc2))


    if maxa -mina <= precision:
        mina, maxa = minmax(mina,maxa,mgrid_a[0],mgrid_a[-1],a_best,precision)
    if maxb - minb <= precision:
        minb, maxb = minmax(minb,maxb,mgrid_b[0],mgrid_b[-1],b_best,precision)
    if maxc - minc <= precision:
        minc, maxc = minmax(minc,maxc,mgrid_c[0],mgrid_c[-1],c_best,precision)

#    print mina, maxa, minb, maxb
    print "min a", mina, "max a", maxa, "min b", minb, "max b", maxb,"min c", minc, "max c", maxc, "best a", a_best, "best b", b_best, "best c", c_best
    return mina, maxa, minb, maxb, minc, maxc, a_best, b_best, c_best, len(abc0)

# hack to prevent  max - min from going below 0.5*precision

def minmax(mina,maxa,MIN,MAX,BEST,PRECISION):
    err = 1e-5
    if ((BEST - 0.5*PRECISION >= MIN) and  (BEST + 0.5*PRECISION <= MAX)):
        mina = BEST - 0.5*PRECISION
        maxa = BEST + 0.5*PRECISION
    else:
        if mina - MIN <= err:
            maxa = MIN + PRECISION
        elif MAX -maxa <= err:
            mina = MAX - PRECISION
        elif BEST - 0.5*PRECISION <= MIN:
            mina = MIN 
            maxa = MIN  + PRECISION
        elif BEST + 0.5*PRECISION >= MAX:
            maxa = MAX
            mina = MAX - PRECISION
        
    return mina, maxa 

